#include <iostream>
#include <fstream>
#include <vector>
#include <cassert>
#include <cmath>

namespace KMeansMultiple {
    std::vector<std::vector<int>>
    get_label_map(int *label_l, int n_cluster, int n_item, int n_overlap) {
        std::vector<std::vector<int>> label_map(n_cluster);
        for (int i = 0; i < n_cluster; i++) {
            label_map[i] = std::vector<int>();
        }
        for (int i = 0; i < n_item; i++) {
            for (int j = 0; j < n_overlap; j++) {
                int cluster = label_l[n_overlap * i + j];
                label_map[cluster].push_back(i);
            }
        }
        return label_map;
    }

    double l2distance(float *a, float *b, int len) {
        double total_distance = 0;
        for (int i = 0; i < len; i++) {
            float diff = a[i] - b[i];
            total_distance += diff * diff;
        }
        double dist = sqrt(total_distance);
        return dist;
    }

    void
    get_weight_l(float *base, int *this_label_l, std::vector<std::vector<int>> &label_map, int dimension, int n_cluster
                 , int item_idx
                 , std::vector<float> &weight_l) {
        int n_overlap = (int) weight_l.size();
        //for each label, calculate the total distance in the cluster
        double *distance_l = new double[n_overlap];
        for (int i = 0; i < n_overlap; i++) {
            int cluster_i = this_label_l[i];
            std::vector<int> cluster_item_l = label_map[cluster_i];
            int cluster_len = (int) cluster_item_l.size();
            float *this_item = base + dimension * item_idx;
            double total_distance = 0;
            for (int j = 0; j < cluster_len; j++) {
                float *iter_item = base + dimension * cluster_item_l[j];
                total_distance += l2distance(this_item, iter_item, dimension);
            }
            distance_l[i] = total_distance / cluster_len;
        }

        //pass through the softmax
        double exp_sum = 0;
        for (int i = 0; i < n_overlap; i++) {
            distance_l[i] = std::exp(-distance_l[i]);
            exp_sum += distance_l[i];
        }
        for (int i = 0; i < n_overlap; i++) {
            weight_l[i] = (float) (distance_l[i] / exp_sum);
        }
    }

}
